#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版文本向量化 + Milvus存储
功能：读取txt文件，向量化，保存到Excel和Milvus数据库
"""

import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
import jieba
import re
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
import uuid
from datetime import datetime


def read_txt_file(file_path, encoding='utf-8'):
    """读取txt文件"""
    try:
        with open(file_path, 'r', encoding=encoding) as file:
            lines = file.readlines()
        lines = [line.strip() for line in lines if line.strip()]
        print(f"成功读取文件 {file_path}，共 {len(lines)} 行文本")
        return lines
    except Exception as e:
        print(f"读取文件失败：{e}")
        return []


def preprocess_text(text):
    """文本预处理"""
    text = re.sub(r'[^\w\s]', ' ', text)
    text = re.sub(r'\d+', ' ', text)
    text = text.lower()
    words = jieba.cut(text)
    text = ' '.join(words)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def vectorize_texts(texts, max_features=100):
    """文本向量化"""
    processed_texts = [preprocess_text(text) for text in texts]
    vectorizer = TfidfVectorizer(max_features=max_features, ngram_range=(1, 2))
    vectors = vectorizer.fit_transform(processed_texts)
    print(f"向量化完成，矩阵形状: {vectors.shape}")
    return vectors.toarray(), vectorizer


def export_to_excel(texts, vectors, output_path):
    """导出到Excel"""
    vector_strings = []
    for vector in vectors:
        vector_str = ','.join([f'{x:.6f}' for x in vector])
        vector_strings.append(vector_str)
    
    data = {
        '原始文本': texts,
        '向量': vector_strings
    }
    
    df = pd.DataFrame(data)
    df.to_excel(output_path, index=False, engine='openpyxl')
    print(f"成功导出到 {output_path}")


def connect_to_milvus(host="localhost", port="19530"):
    """连接到Milvus数据库服务器"""
    try:
        # 连接到Docker部署的Milvus服务器
        connections.connect("default", host=host, port=port)
        print(f"成功连接到Milvus服务器 {host}:{port}")
        return True
    except Exception as e:
        print(f"连接Milvus服务器失败: {e}")
        print("请确保Milvus Docker容器正在运行")
        print("检查命令: docker ps | grep milvus")
        return False


def create_milvus_collection(collection_name, vector_dim):
    """创建Milvus集合"""
    # 如果集合已存在，先删除
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)
        print(f"删除已存在的集合: {collection_name}")

    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=5000),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
        FieldSchema(name="timestamp", dtype=DataType.VARCHAR, max_length=50)
    ]

    # 创建集合
    schema = CollectionSchema(fields, f"文本向量集合")
    collection = Collection(collection_name, schema)
    print(f"成功创建集合: {collection_name}")
    
    # 创建索引
    index_params = {
        "metric_type": "L2",
        "index_type": "IVF_FLAT", 
        "params": {"nlist": 128}
    }
    collection.create_index("vector", index_params)
    print("成功创建向量索引")
    
    return collection


def insert_to_milvus(collection, texts, vectors):
    """插入数据到Milvus"""
    # 准备数据
    ids = [str(uuid.uuid4()) for _ in range(len(texts))]
    timestamps = [datetime.now().strftime("%Y-%m-%d %H:%M:%S") for _ in range(len(texts))]
    
    # 插入数据
    data = [ids, texts, vectors.tolist(), timestamps]
    mr = collection.insert(data)
    collection.flush()
    
    print(f"成功插入 {len(texts)} 条向量数据到Milvus")
    return mr.primary_keys


def search_similar_texts(collection, vectorizer, query_text, top_k=5):
    """搜索相似文本"""
    # 加载集合到内存
    collection.load()
    
    # 向量化查询文本
    processed_query = preprocess_text(query_text)
    query_vector = vectorizer.transform([processed_query]).toarray()
    
    # 搜索参数
    search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
    
    # 执行搜索
    results = collection.search(
        query_vector.tolist(),
        "vector",
        search_params,
        limit=top_k,
        output_fields=["text", "timestamp"]
    )
    
    # 处理结果
    similar_texts = []
    for hits in results:
        for hit in hits:
            similar_texts.append({
                "text": hit.entity.get("text"),
                "distance": hit.distance,
                "timestamp": hit.entity.get("timestamp")
            })
    
    return similar_texts


def process_txt_to_milvus(input_file, output_excel=None, collection_name="text_vectors",
                         max_features=100, milvus_host="localhost", milvus_port="19530"):
    """完整处理流程：txt -> 向量化 -> Excel + Milvus"""
    print("=" * 50)
    print("开始文本向量化和Milvus存储")
    print(f"Milvus服务器: {milvus_host}:{milvus_port}")
    print("=" * 50)

    # 1. 读取文件
    texts = read_txt_file(input_file)
    if not texts:
        print("没有读取到有效文本，程序退出")
        return None, None, None

    # 2. 向量化
    vectors, vectorizer = vectorize_texts(texts, max_features)

    # 3. 导出到Excel（可选）
    if output_excel:
        export_to_excel(texts, vectors, output_excel)

    # 4. 存储到Milvus
    if connect_to_milvus(milvus_host, milvus_port):
        collection = create_milvus_collection(collection_name, vectors.shape[1])
        insert_to_milvus(collection, texts, vectors)

        print("=" * 50)
        print("处理完成！数据已保存到Milvus服务器")
        print(f"集合名称: {collection_name}")
        print(f"数据量: {len(texts)} 条")
        print("=" * 50)

        return texts, vectors, vectorizer, collection
    else:
        print("Milvus连接失败，仅保存到Excel")
        return texts, vectors, vectorizer, None


def demo_search(collection, vectorizer):
    """演示搜索功能"""
    if collection is None or vectorizer is None:
        print("搜索功能不可用")
        return
    
    print("\n" + "=" * 50)
    print("演示相似性搜索:")
    
    query = "文本向量化"
    similar_texts = search_similar_texts(collection, vectorizer, query, top_k=3)
    
    print(f"查询文本: {query}")
    print("相似文本:")
    for i, result in enumerate(similar_texts, 1):
        print(f"{i}. {result['text'][:50]}... (距离: {result['distance']:.4f})")


if __name__ == "__main__":
    # 配置参数
    input_file = "sample.txt"                    # 输入txt文件
    output_excel = "向量结果_docker_milvus.xlsx"  # 输出Excel文件（可选）
    collection_name = "my_text_vectors"          # Milvus集合名称
    max_features = 50                           # 向量维度

    # Milvus Docker服务器配置
    milvus_host = "localhost"                   # Milvus服务器地址
    milvus_port = "19530"                       # Milvus服务器端口

    print("配置信息:")
    print(f"- 输入文件: {input_file}")
    print(f"- 输出Excel: {output_excel}")
    print(f"- Milvus服务器: {milvus_host}:{milvus_port}")
    print(f"- 集合名称: {collection_name}")
    print(f"- 向量维度: {max_features}")

    # 处理文件
    texts, vectors, vectorizer, collection = process_txt_to_milvus(
        input_file=input_file,
        output_excel=output_excel,
        collection_name=collection_name,
        max_features=max_features,
        milvus_host=milvus_host,
        milvus_port=milvus_port
    )

    # 演示搜索
    if collection:
        demo_search(collection, vectorizer)
